#!/usr/bin/python2.4

# pylint: disable-msg=C6304

"""Reads tidal current prediction images from DFO."""

import glob
import math
import os
import re
import sys
import time

import imagereader
import matrix
from matrix import DM
from matrix import LatLon
from matrix import Point
from matrix import Polar
from matrix import Ratio
from matrix import TransformStack
import timeformat
import vectorfield


class DFOVectorField(vectorfield.VectorField):

  def __init__(self, name='unknown'):
    super(DFOVectorField, self).__init__(name)
    # These deltas are in MINUTES.
    self.SetLatLonFuzzMinutes(4, 4)
    self.SetTimeFuzzSeconds(1 * 3600)


def GetSummaryCSVFileName():
  """Return the name of the file holding the current tidal vector field."""
  return 'dfodata.csv'


class WrappedDFOPixMatrix(object):
  """Provide a class implementing SetPixel to wrap around a DFO pix matrix."""

  def __init__(self, pix):
    self._pix = pix

  def GetPixelRaw(self, point):
    if point.x < 0 or point.y < 0:
      return None
    if point.x >= self._clipped_height_y or point.y >= self._clipped_width_x:
      return None
    pixel = self._pix[point.x, point.y]
    return pixel

  def SetPixel(self, point, color):
    try:
      self._pix[point.x, point.y] = color
    except (IndexError, AttributeError), unused_e:
      pass


class DFOImageReader(imagereader.ImageReader):
  """Reads PPM image files from DFO holding forecast tidal current data."""

  def __init__(self, tesseract_api, trace_level=0):
    super(DFOImageReader, self).__init__(tesseract_api, trace_level)
    self._is_calibration = False

  def GetSourceFiles(self, desired):
    """Get information on gif or ppm image files to read.

    Args:
      desired: An unordered collection of source file names.

    Returns:
      A list of all file names to read.
    """
    if 'today' in desired or 'now' in desired:
      image_dir = 'dfo/2011-0[56]-[012][0-9]'  # When this software was developed
        # TODO(rrh): Figure out why we can't read and process GIF files.
      file_suffix = 'pr_*_jdf.ppm'

    if 'oor' in desired or 'oos' in desired:
      image_dir = 'dfo/2011-05-1*.oor'  # for Oregon Offshore
      file_suffix = 'pr_*_oos.ppm'

    if 'ss' in desired or 'swift' in desired:
      image_dir = 'dfo/2011-05-[23]*.swift'  # for Swiftsure
      file_suffix = 'pr_*_swift.ppm'

    if 'oor' in image_dir or 'swift' in image_dir:
      # Special product forecasts, used for oor or swiftsure.
      self._full_width_x = 720
      self._full_height_y = 504
      self._left_edge = -DM(125, 06)
      self._right_edge = -DM(123, 12)
      self._upper_edge = DM(48, 48)
      self._lower_edge = DM(48, 06)
    else:
      # Daily forecasts
      self._full_width_x = 864
      self._full_height_y = 504
      self._left_edge = -DM(124, 54)
      self._right_edge = -DM(123, 00)
      self._upper_edge = DM(48, 42)
      self._lower_edge = DM(48, 06)

    return glob.glob(
         os.path.join(imagereader.WORKDIR, image_dir, file_suffix))

  def ImageBoundingBox(self):
    """Return a pixel bounding box for the pixels of interest in the image."""
    # TODO(rrh): Crop out and save just the useful image
    # These numbers are empirical.

    left = 30  # discard white space on left of image
    lower = self._full_height_y - 25  # chops off the long labels on bottom
    right = self._full_width_x - 75  # chops off the latitude labels on right
    if self._is_calibration:
      upper = 0
    else:
      upper = 75  # chops off the key and date across the top
    self._clipped_width_x = abs(right - left)
    self._clipped_height_y = abs(upper - lower)

    crop_box = (left, upper, right, lower)
    return crop_box

  def TextBoundingBox(self):
    """Return a pixel bounding box holding the date/time text of interest."""
    if self._full_height_y == 504 and self._full_width_x == 864:
      return (150, 10, 350, 40)
    if self._full_height_y == 504 and self._full_width_x == 720:
      return (150-50, 10, 350-50, 40)
    else:
      raise Exception('Unknown text bounding box')

  def GetImageTime(self, string_scanned, file_name):
    """Convert the string_scanned from OCR into date and time.

    Sets self._unix_epoch_time as a side effect.
    TODO(rrh): should return the unix_epoch_time as well.

    Args:
      string_scanned: the string that was returned by the OCR engine.
      file_name: the name of the image file from which the string was extracted.

    Returns:
      The string that was scanned.
    """
    if True:
      timestamp_filename = self.MakeTimeFromFileName(file_name)
      timestamp = timestamp_filename  # Times via filename are 100% accurate
    else:
      timestamp_ocr = self.MakeTimeFromOCR(string_scanned)
      timestamp = timestamp_ocr  # Flaky, but better to manage

    if timestamp:
      self._unix_epoch_time = timestamp
      print 'file_name=%s time_extracted=%s' % (
          file_name,
          timeformat.FormatTime(self._unix_epoch_time))
      return string_scanned
    else:
      print 'file_name=%s string_scanned=%18s ==> MISS' % (
          file_name,
          string_scanned)
      return None

  def MakeTimeFromOCR(self, string_scanned):
    """Converts an OCR string from Tesseract into a unix time stamp.

    The (untuned) Tesseract OCR is unreliable,
    and I could not really make this work.
    Instead, use the knowledge in the file names.

    Args:
      string_scanned: The string from the OCR engine.

    Returns:
      The time since the unix epoch, seconds.
    """
    string_scanned = string_scanned.strip()
    string_scanned = string_scanned.replace('\xe2\x80\x98', '1')
    string_scanned = re.sub(r'^ *Z', '2', string_scanned)
    string_scanned = re.sub(r'E', '8', string_scanned)
    string_scanned = re.sub(r'T', '1', string_scanned)
    string_scanned = re.sub(r' M ', ' 14 ', string_scanned)

    # For the DFO data,
    # TODO(rrh): Get H for 11
    # TODO(rrh): Get H for 14
    # TODO(rrh): Get 15 for 18(!)
    # Need more training data for Tesseract?

    d = r'[0-9BUDOI]'
    d4 = r'(' + d + r'{4}' + r')'
    d2 = r'(' + d + r'{2}' + r')'
    re_string = (
        r'^' + d4
        + r' *' + d2  # Sometimes get 0 spaces after the year
        + r' *' + d2
        + r' *' + d2
        # + r'Z$'
        )

    rex = re.compile(re_string)
    m = rex.match(string_scanned)
    if m:
      year = int(imagereader.UnTesseract(m.group(1)))
      month = int(imagereader.UnTesseract(m.group(2)))
      date = int(imagereader.UnTesseract(m.group(3)))
      hour = int(imagereader.UnTesseract(m.group(4)))
      while hour >= 24:
        hour -= 24
        date += 1
      reformatted = '%d-%02d-%02d %02d:00 %s' % (
          year,
          month,
          date,
          hour,
          'PDT'  # aka localtime
          )
      time_fmt = '%Y-%m-%d %H:%M %Z'
      timestamp = time.mktime(time.strptime(reformatted, time_fmt))
      return timestamp
    else:
      return None

  def MakeTimeFromFileName(self, file_name):
    """Make a valid time from the file name.

    Args:
      file_name: The name of the file holding the image.

    Returns:
      The number of seconds since the unix epoch.
    """
    # Now, determine the local time from the file name.
    # The file names encode dates/times in UTC.
    # We downloaded the files into directories named %Y-%m-%d
    # and files come to us in the form pr_%h_*.gif
    basename = os.path.basename(file_name)
    dirname = os.path.basename(os.path.dirname(file_name))
    m = re.match('(....)-(..)-(..)', dirname)
    utc_diryear = m.group(1)
    utc_dirmonth = m.group(2)
    utc_dirdate = m.group(3)
    utc_filetime = basename[3:5]  # values 01..24
    reformatted = '%d-%02d-%02d %02d:00 %s' % (
        int(utc_diryear),
        int(utc_dirmonth),
        int(utc_dirdate),
        int(utc_filetime)-1,  # compensate for 01..24 in file names
        'UTC',
        )
    time_fmt = '%Y-%m-%d %H:%M %Z'
    alt_time = time.mktime(time.strptime(reformatted, time_fmt))

    alt_time += 1 * 3600  # re-compenasate for 01..24 in file names
    alt_time -= 1 * 3600  # Gross hack: compensate timezone in Daylight time
    alt_time -= 7 * 3600  # Gross hack: then compensate for PDT from UTC

    if False:
      print '%s reformatted={%s} alt_time={%s}' % (
          file_name, reformatted, timeformat.FormatTime(alt_time))
    return alt_time

  def SelectPreferredColor(self, file_name, image):
    """Select from image a new image containing just the preferred colors.

    For DFO, we just retain pixels with a tinge of green to them; these
    make up the current vectors.

    Args:
      file_name: The name of the file we read the image from.
      image: An image.

    Returns:
      A new image.
    """

    if self._is_calibration:
      green_image = SelectPureBlack(image)
    else:
      green_image = SelectAnyGreen(image)

    return green_image

  def CustomAnalysis(self, file_name, output_file=sys.stdout):
    """Do the custom analysis for the DFO images.

    Args:
      file_name: The name of the image file that we're working on.
      output_file: A file to write the output CSV to.
    """
    # Apparently, we have to copy() the image before we can load() it.
    trace_level = 1
    if trace_level >= 1:
      print 'Reading file %s' % (file_name,)
    image = self._images[file_name].copy()
    pix = image.load()

    wrapped_pix = WrappedDFOPixMatrix(pix)

    # Look for all Green points (eg, those not pure white)
    green_set = set()
    for y in xrange(0, self._clipped_height_y):
      for x in xrange(0, self._clipped_width_x):
        p = Point(x, y)
        pixvalue = pix[p.x, p.y]
        if pixvalue != imagereader.WHITE:
          if self._is_calibration:
            if pixvalue == imagereader.BLACK:
              green_set.add(p)
          else:
            if pixvalue != imagereader.BLACK:
              green_set.add(p)

    dither_range = 2  # how far from the starting point we look for adjacents
    visited_set = set()
    arrow_number = 0
    arrow_sets = []
    for starter_point in green_set:
      if starter_point in visited_set:
        continue

      arrow_set = set()
      arrow_set.add(starter_point)

      # Work on points that we have not yet visited but are nearby the seeds.
      points_to_visit = set()
      points_to_visit.add(starter_point)
      while points_to_visit:
        new_points_to_visit = set()
        for p in points_to_visit:
          for x in xrange(-dither_range, dither_range+1):
            for y in xrange(-dither_range, dither_range+1):
              probe = Point(p.x + x, p.y + y)
              if probe in visited_set:
                continue
              if probe in green_set and probe not in arrow_set:
                arrow_set.add(probe)
                new_points_to_visit.add(probe)  # Extend the frontier to explore
        points_to_visit = new_points_to_visit
      visited_set |= arrow_set

      # TODO(rrh): some arrow_set instances may be enormous (68K members)
      # (This turned out to be due to a poorly sized box.)
      if len(arrow_set) >= 2 and len(arrow_set) <= 50:
        if trace_level > 1:
          print 'arrow_set for arrow %3d ==> %3d' % (
              arrow_number, len(arrow_set),)
        arrow_number += 1
        arrow_sets.append(arrow_set)

    # Determine the average point (centroid) for each arrow;
    # in the next step we'll use that as the sort key so that
    # we peel off arrows in some sorted order.
    avg_point_to_arrow_set = {}
    for arrow_set in arrow_sets:
      # Find the average point.
      sigma_x = 0
      sigma_y = 0
      n = 0
      for p in arrow_set:
        sigma_x += p.x
        sigma_y += p.y
        n += 1
      avg_point = Point(sigma_x/n, sigma_y/n)
      avg_point_to_arrow_set[avg_point] = arrow_set

    # Iterate in sorted order by avg_point.
    for avg_point in sorted(avg_point_to_arrow_set):
      arrow_set = avg_point_to_arrow_set[avg_point]

      # Find the 2 points that are furthest apart, to determine length.
      max_distance = 0
      best_p0 = None
      best_p1 = None
      for p0 in arrow_set:
        for p1 in arrow_set:
          distance = p0.DistanceFrom(p1)
          if distance > max_distance:
            max_distance = distance
            best_p0 = p0
            best_p1 = p1

      # Count the number of adjacent points to the 2 points furthest apart
      adjacents_p0 = 0
      adjacents_p1 = 0
      adjacent_delta = 1

      best_g0 = 0  # value of lightest green pixel near p0
      best_g1 = 0  # value of lightest green pixel near p1

      for dx in xrange(-adjacent_delta, adjacent_delta+1):
        for dy in xrange(-adjacent_delta, adjacent_delta+1):
          probe = Point(best_p0.x + dx, best_p0.y + dy)
          if probe in arrow_set:
            adjacents_p0 += 1
            g = imagereader.GetGreen(pix[probe.x, probe.y])
            if g > best_g0:
              best_g0 = g

          probe = Point(best_p1.x + dx, best_p1.y + dy)
          if probe in arrow_set:
            adjacents_p1 += 1
            g = imagereader.GetGreen(pix[probe.x, probe.y])
            if g > best_g1:
              best_g1 = g

      # The head of the arrow has a pixel within distance 1 that has more
      # white in it
      if (
          (adjacents_p0 == adjacents_p1) or
          (adjacents_p0 <= 2 and adjacents_p1 <= 2)
          ):
        if best_g1 > best_g0:
          p_from = best_p0
          p_to = best_p1
        else:
          p_from = best_p1
          p_to = best_p0

      elif True:
        # The head of the arrow has more points adjacent to it than the tail.
        if adjacents_p0 > adjacents_p1:
          p_from = best_p1
          p_to = best_p0
        else:
          p_from = best_p0
          p_to = best_p1
      else:
        # The head of the arrow is closest to the average point.
        # This does not do so well for fat short arrows.
        if best_p0.DistanceFrom(avg_point) > best_p1.DistanceFrom(avg_point):
          p_from = best_p0
          p_to = best_p1
        else:
          p_from = best_p1
          p_to = best_p0

      display_radians = math.atan2(p_to.y - p_from.y, p_to.x - p_from.x)
      p_center = Point((p_from.x + p_to.x)/2, (p_from.y + p_to.y)/2)
      stack = TransformStack()
      stack.Translate(p_center)
      stack.Rotate(display_radians)
      stack.Scale(max_distance, max_distance)
      dy = 0
      dy = 0.25  # offset arrow from green blobs underneath

      # Draw shaft of arrow in RED
      matrix.BresenhamLineFloats(wrapped_pix,
                                 stack.Transform(Point(-1.0, dy)),
                                 stack.Transform(Point(1.0, dy)),
                                 imagereader.RED)

      # Draw barb of arrow in blue, on top of the shaft.
      matrix.BresenhamLineFloats(wrapped_pix,
                                 stack.Transform(Point(1.0, dy)),
                                 stack.Transform(Point(1.0-0.3, dy+0.3)),
                                 imagereader.BLUE)

      # 11 pixels that we extract on the calibration line
      # represents 1 meter/second = 1.94384 knots
      speed_knots = (max_distance / 11.0) * 1.94384
      # TODO(rrh): Do what it takes to understand this.
      # Here, we have the positions of x and y interchanged to have
      # current trending North come out as 0 degrees
      theta_radians = math.atan2(p_to.x - p_from.x, p_to.y - p_from.y)
      theta_degrees = 180 - math.degrees(theta_radians)

      center_latlon = self.ConvertPixelToLatLon(p_center)
      coords = vectorfield.CoordsToVector(
          self._unix_epoch_time,
          center_latlon,
          Polar(speed_knots, theta_degrees),
          attribution=file_name)
      print >>output_file, '%s' % (coords.CSVFormat(),)

    if True:
      # Save the extracted and overdrawn image to help eyeball conversion bugs.
      output_fname = os.path.join(
          os.path.dirname(file_name),
          'extracted.' + os.path.basename(file_name))
      image.save(output_fname, 'PPM')

  def ConvertPixelToLatLon(self, pixel_point):
    lat = Ratio(0, pixel_point.y, self._clipped_height_y,
                self._upper_edge, self._lower_edge)
    lon = Ratio(0, pixel_point.x, self._clipped_width_x,
                self._left_edge, self._right_edge)
    latlon = LatLon(lat, lon)
    return latlon


def SelectAnyGreen(base_image):
  """Modifies base_image to discard non-green pixels.

  Used to select the current vectors out of the DFO images.

  Args:
    base_image: The base image to discard non-green points.

  Returns:
    base_image modified to retain just interesting green pixels.
  """
  green_data = []
  for p in base_image.getdata():
    # print '%s' % (p,)
    if p[0] > 240 and p[1] > 240 and p[2] > 240:
      green_data.append(imagereader.WHITE)  # Mostly white to pure white
    elif p[0] < 10 and p[1] < 10 and p[2] < 10:
      green_data.append(imagereader.BLACK)
    elif p[1] > 0 and p[2] > 0 and p[1] > p[2]:  # some yellow tinge
      green_data.append((0, p[1], p[2]))
    else:
      green_data.append(imagereader.WHITE)
  base_image.putdata(green_data)
  return base_image


def SelectPureBlack(base_image):
  """Modifies base_image to discard non-pure-black pixels.

  Used to select the calibration(key) vector out of the DFO images.

  Args:
    base_image: The base image to discard non-pure black points.

  Returns:
    base_image modified to retain just interesting pure black pixels.
  """
  green_data = []
  for p in base_image.getdata():
    if p == imagereader.BLACK:
      green_data.append(imagereader.BLACK)
    else:
      green_data.append(imagereader.WHITE)
  base_image.putdata(green_data)
  return base_image


def ReadGIFFromDFO(desired):
  """Read all of the .gif files with tidal forecasts from the DFO server."""
  dfo_http = 'http://www.pac.dfo-mpo.gc.ca/sci/juandefuca/'

  if 'today' in desired or 'now' in desired:
    # DFO produced today's and tomorrow's tidal prediction images
    imagereader.GetGIFFromServer(
        dfo_http + 'images/today/',
        os.path.join(
            imagereader.WORKDIR,
            'dfo',
            time.strftime('%Y-%m-%d', time.localtime(time.time()))),
        'pr_%02d_jdf.gif', 24)

    imagereader.GetGIFFromServer(
        dfo_http + 'images/tomorrow/',
        os.path.join(
            imagereader.WORKDIR,
            'dfo',
            time.strftime('%Y-%m-%d', time.localtime(time.time() + 24 * 3600))),
        'pr_%02d_jdf.gif', 24)

  if 'oor' in desired or 'oos' in desired:
    # DFO produced "special product" tidal predictions for Oregon Offshore Race
    # All times in GMT
    # This loop specalized for dates of the Oregon Offshore 2011
    for date in [12, 13, 14, 15]:
      imagereader.GetGIFFromServer(
          dfo_http + 'images/oos/201105%02d/' % (date,),
          os.path.join(
              imagereader.WORKDIR,
              'dfo',
              '2011-05-%02d.oor' % (date,)),
          'pr_%02d_oos.gif',
          24)

  if 'ss' in desired or 'swift' in desired:
    # DFO produced "special product" tidal predictions for Swiftsure
    for date in [27, 28, 29]:
      imagereader.GetGIFFromServer(
          dfo_http + 'images/swiftsure/201105%02d/' % (date,),
          os.path.join(
              imagereader.WORKDIR,
              'dfo',
              '2011-05-%02d.swift' % (date,)),
          'pr_%02d_swift.gif', 24)
